# ===============================================================
# 文件名: CT_create.py
# 功能: 机动目标轨迹数据生成器（真实轨迹+观测数据）
# 侦测目标距离范围: 926m~37040m
# 侦测目标速度范围: -340m/s~340m/s
# CT模型中转弯速率: -10°/s~10°/s
# 改变间隙 0.1°/s
# ===============================================================
import numpy as np
import random as rd
import math

from filterpy.kalman import UnscentedKalmanFilter
from filterpy.kalman import JulierSigmaPoints

import matplotlib.pyplot as plt
from scipy.io import savemat

import time


ST = 0.1  # 采样间隔
DATA_LENGTH = 50  # 采样点数-->采样时间

VELOCITY_MAX = 340  # 最大速度
VELOCITY_MIN = -340  # 最小速度

DISTANCE_MAX = 37040  # 最大距离
DISTANCE_MIN = 926  # 最小距离

TURN_RATE_MAX = 10  # 最大转弯率
TURN_RATE_MIN = -10  # 最小转弯率
TURN_RATE_INTERVAL = 0.1  # 转弯率间隙

DIM_STATE = 5  # 状态维数
DIM_OBS = 2  # 观测维数

DEFAULT_SENSOR_X = 0.0  # 传感器位置X坐标
DEFAULT_SENSOR_Y = 0.0  # 传感器位置Y坐标

INI_POS_NOISE = 30  # 初始位置噪声
INI_VEL_NOISE = 3  # 初始速度噪声
def time_decorator(func):
    def wrapper(*args, **kwargs):
        start_time = time.perf_counter()
        result = func(*args, **kwargs)
        end_time = time.perf_counter()
        elapsed_time = end_time - start_time
        print(f"函数 {func.__name__} 耗时: {elapsed_time:.4f}秒")
        return result
    return wrapper


def normalize_angle(angle):
    """将角度归一化到[-pi, pi]范围内"""
    while angle > np.pi:
        angle -= 2 * np.pi
    while angle < -np.pi:
        angle += 2 * np.pi
    return angle


def simulate_ct_step(x, dt):
    """
    根据CT模型计算下一个时间步的真实状态
    状态向量x = [px, py, v, yaw, omega]
    dt: 时间步长
    """
    px, py, v, yaw, omega = x
    # 避免 omega除零错误
    if abs(omega) < 1e-6:  # 接近零转弯率，近似CV模型
        px_next = px + v * np.cos(yaw) * dt
        py_next = py + v * np.sin(yaw) * dt
    else:  # 标准CT模型更新
        v_over_omega = v / omega
        # 计算下一时刻的航向角（用于位置更新）
        yaw_next_calc = yaw + omega * dt
        px_next = px + v_over_omega * (np.sin(yaw_next_calc) - np.sin(yaw))
        py_next = py + v_over_omega * (-np.cos(yaw_next_calc) + np.cos(yaw))

    # 更新航向角、速率、转弯率（在理想CT模型中，v和omega不变）
    yaw_next = normalize_angle(yaw + omega * dt)  # 确保航向角在[-pi, pi]范围内
    v_next = v  # 速度保持不变
    omega_next = omega  # 转弯率保持不变

    return np.array([px_next, py_next, v_next, yaw_next, omega_next])


def create_true_orbit(turn_rate, data_len=DATA_LENGTH, dim_state=DIM_STATE):
    """
    生成一条真实轨迹
    :param turn_rate: 轨迹转弯率
    :param data_len: 轨迹长度
    :param dim_state: 状态维数，默认5维
    :return: 真实轨迹[px,py,v,yaw,omega]
    px: 目标位置X坐标（单位：米）
    py: 目标位置Y坐标（单位：米）
    v: 目标速度（标量，单位：米/秒）
    yaw: 目标航向角（弧度）
    omega: 目标转弯率（弧度/秒）
    """
    # 产生一个data_len*dim_state的零矩阵，用于存储真实轨迹
    data = np.zeros((data_len, dim_state), dtype='float64')
    distance_change_max = VELOCITY_MAX * ST * data_len

    # 轨迹起点距离
    start_point_distance = rd.uniform(DISTANCE_MIN + distance_change_max, DISTANCE_MAX - distance_change_max)

    # 轨迹起点的笛卡尔坐标（二维）
    start_point_direction = (rd.random() - 0.5) * 2 * np.pi  # 生成一个(-180°, 180°)范围内的均匀分布的角度
    distance_x = start_point_distance * np.cos(start_point_direction)  # X方向位置
    distance_y = start_point_distance * np.sin(start_point_direction)  # Y方向位置

    # 轨迹起点的速度
    start_point_velocity = rd.uniform(VELOCITY_MIN, VELOCITY_MAX)
    start_point_velocity_direction = (rd.random() - 0.5) * 2 * np.pi
    velocity_x = start_point_velocity * np.cos(start_point_velocity_direction)   # Target X direction velocity
    velocity_y = start_point_velocity * np.sin(start_point_velocity_direction)   # Target Y direction velocity

    initial_state = np.array([distance_x, distance_y, np.abs(start_point_velocity), start_point_velocity_direction, np.deg2rad(turn_rate)], dtype='float64')

    for i in range(data_len):
        data[i, :] = initial_state
        initial_state = simulate_ct_step(initial_state, ST)

    return data


def create_true_orbit_batch(batch_size, data_len=DATA_LENGTH, dim_state=DIM_STATE):
    """
    生成一批真实轨迹
    :param batch_size: 批量大小
    :param data_len: 轨迹长度
    :param dim_state: 状态维数，默认5维
    :return: 真实轨迹[px,py,v,yaw,omega]
    """
    # 产生一个batch_size*data_len*dim_state的零矩阵，用于存储真实轨迹
    data = np.zeros((batch_size, data_len, dim_state), dtype='float64')

    for i in range(batch_size):
        # 随机生成转弯率
        min_bound_index = int(TURN_RATE_MIN / TURN_RATE_INTERVAL)
        max_bound_index = int(TURN_RATE_MAX / TURN_RATE_INTERVAL)
        turn_rate = rd.randint(min_bound_index, max_bound_index) * TURN_RATE_INTERVAL
        data[i] = create_true_orbit(turn_rate, data_len, dim_state)
    return data


def hx_direct_position_ct(x, sensor_x=DEFAULT_SENSOR_X, sensor_y=DEFAULT_SENSOR_Y):
    """
    观测函数：直接测量目标位置（px,py）
    适用于CT模型状态：[px, py, v, yaw, omega]
    返回观测向量：[measured_px, measured_py]
    """
    px, py, _, _, _ = x  # 只关心状态中的位置分量
    # 传感器直接测量位置，与传感器自身位置无关
    return np.array([px, py])


def hx_distance_azimuth_ct(x, sensor_x=DEFAULT_SENSOR_X, sensor_y=DEFAULT_SENSOR_Y):
    """
    观测函数：测量目标的距离和方位角
    适用于CT模型状态：[px, py, v, yaw, omega]
    返回观测向量：[distance, azimuth]
    """
    px, py, _, _, _ = x # 只关心状态中的位置分量
    # 计算相对位置
    relative_x = px - sensor_x
    relative_y = py - sensor_y

    # 计算距离和方位角
    distance = np.sqrt(relative_x**2 + relative_y**2)
    azimuth = np.arctan2(relative_y, relative_x)

    return np.array([distance, azimuth])


def hx_distance_azimuth_velocity_ct(x, sensor_x=DEFAULT_SENSOR_X, sensor_y=DEFAULT_SENSOR_Y):
    """
    观测函数：测量目标的距离、方位角和速度
    适用于CT模型状态：[px, py, v, yaw, omega]
    返回观测向量：[distance, azimuth, velocity]
    """
    px, py, v, _, _ = x  # 只关心状态中的位置和速度分量
    # 计算相对位置
    relative_x = px - sensor_x
    relative_y = py - sensor_y

    # 计算距离和方位角
    distance = np.sqrt(relative_x**2 + relative_y**2)
    azimuth = np.arctan2(relative_y, relative_x)
    azimuth = normalize_angle(azimuth)  # 归一化方位角到[-pi, pi]范围内

    return np.array([distance, azimuth, v])


def create_mea_orbit_batch(batch_size, data_len=DATA_LENGTH, dim_state=DIM_STATE, dim_obs=DIM_OBS):
    """
    生成一批观测轨迹
    :param batch_size: 批量大小
    :param data_len: 轨迹长度
    :param dim_state: 状态维数，默认5维
    :param dim_obs: 观测维数，默认2维
    :return:
    真实轨迹[px, py, v, yaw, omega]
    观测轨迹[distance, azimuth]
    """

    obs_data = np.zeros((batch_size, data_len, dim_obs), dtype='float64')

    true_data = create_true_orbit_batch(batch_size, data_len, dim_state)

    for i in range(batch_size):
        distance_noise = rd.uniform(8.0, 13.0)  # distance noise
        distance_var = np.square(distance_noise)  # distance noise variance
        azimuth_noise = rd.uniform(7.0, 9.0)  # azimuth noise
        azimuth_var = np.square(azimuth_noise / 1000)  # azimuth noise variance

        obs_data[i, :, 0] = np.sqrt(np.square(true_data[i, :, 0]) +
                                    np.square(true_data[i, :, 1])
                                    ) + np.random.normal(0, distance_var, data_len)  # 观测值距离加上误差

        obs_data[i, :, 1] = (np.arctan2(true_data[i, :, 1], true_data[i, :, 0]) +
                             np.random.normal(0, azimuth_var, data_len)
                             )  # 观测值方位角加上误差
    return true_data, obs_data


@time_decorator
def create_ukf_orbit_batch(batch_size, data_len=DATA_LENGTH, dim_state=DIM_STATE, dim_obs=DIM_OBS,
                           ini_pos_noise=INI_POS_NOISE, ini_vel_noise=INI_VEL_NOISE):
    """
    生成一批UKF估计轨迹
    :param batch_size: 批量大小
    :param data_len: 轨迹长度
    :param dim_state: 状态维数，默认5维
    :param dim_obs: 观测维数，默认2维
    :param ini_pos_noise: 初始位置噪声标准差
    :param ini_vel_noise: 初始速度噪声标准差
    :return:
    true_data：[px, py, v, yaw, omega]
    obs_data：[distance, azimuth]
    tracking_data：[px,py, v, yaw, omega]
    error_data: [px, py, v, yaw, omega]
    """

    true_data, obs_data = create_mea_orbit_batch(batch_size=256, data_len=DATA_LENGTH, dim_state=DIM_STATE, dim_obs=DIM_OBS)
    tracking_data = np.zeros((batch_size, data_len, dim_state), dtype='float64')
    residuals_history_all = np.zeros((batch_size, data_len-1, dim_obs), dtype='float64')  # 用于存储所有批次的残差历史

    for i in range(batch_size):

        distance_noise = rd.uniform(80.0, 130.0)  # distance noise
        distance_var = np.square(distance_noise)  # distance noise variance
        azimuth_noise = rd.uniform(7.0, 9.0)  # azimuth noise
        azimuth_var = np.square(azimuth_noise / 1000)  # azimuth noise variance
        R = np.diag([distance_var, azimuth_var])  # 观测噪声协方差矩阵

        accelerate_noise = rd.uniform(8.0, 13.0)  # 加速度噪声
        position_noise = accelerate_noise * 0.5 * np.square(ST) * 10  # 位置噪声
        velocity_noise = accelerate_noise * ST * 5  # 速度噪声
        Q = np.diag([position_noise**2, position_noise**2, velocity_noise**2, 0.001, 0])    # 过程噪声协方差矩阵

        my_sigma_points = JulierSigmaPoints(n=dim_state, kappa=3-dim_state)  # 创建 Sigma Points 对象
        my_ukf = UnscentedKalmanFilter(dim_x=dim_state,
                                       dim_z=dim_obs,
                                       dt=ST,
                                       hx=hx_distance_azimuth_ct,
                                       fx=simulate_ct_step,
                                       points=my_sigma_points)  # 创建 UKF 实例
        my_ukf.x = true_data[i, 0, :]
        my_ukf.P = np.diag([0.001**2, 0.001**2, 0.001**2, 0.001**2, 0.001**2])  # 初始状态协方差
        my_ukf.Q = Q
        my_ukf.R *= R

        ukf_states = []
        residuals_history = []  # 用于存储残差历史
        ukf_states.append(my_ukf.x.copy())  # 存储初始估计

        for j in range(data_len-1):
            my_ukf.predict()
            my_ukf.update(obs_data[i, j + 1, :])  # 更新状态

            current_residual = my_ukf.y.copy()  # 获取本次更新的残差
            residuals_history.append(current_residual)  # 存储残差

            ukf_states.append(my_ukf.x.copy())  # 存储每个时间步的估计状态

        ukf_states = np.array(ukf_states)
        residuals_history = np.array(residuals_history)  # 转换为NumPy数组

        tracking_data[i] = ukf_states
        residuals_history_all[i] = residuals_history  # 存储所有批次的残差历史

    return true_data, obs_data, tracking_data, residuals_history_all


def my_plot(true_data, tracking_data, residual_data):
    """
    绘制轨迹
    :param true_data: 真实轨迹数据[batch_size, data_len, dim_state]
    :param tracking_data: 估计轨迹数据[batch_size, data_len, dim_state]
    :param residual_data: 残差数据[batch_size, data_len-1, dim_obs]
    :return: None
    """
    for i in range(10):
        plt.figure(figsize=(10, 8))
        plt.plot(true_data[i, :, 0], true_data[i, :, 1], linestyle='-')
        plt.plot(tracking_data[i, :, 0], tracking_data[i, :, 1], linestyle='--')
        plt.title(f'true_orbit & UKF_tracking_orbit-----{i+1}')  # 标题
        plt.xlabel('X (m)')
        plt.ylabel('Y (m)')
        plt.legend(['true_orbit', 'UKF_tracking_orbit'])  # 图例
        plt.grid(True)
        plt.show()

        plt.figure(figsize=(10, 8))
        time_axis_residuals = np.arange(1, DATA_LENGTH) * ST

        plt.subplot(2, 1, 1)
        plt.plot(time_axis_residuals, residual_data[i, :, 0], linestyle='-', label='Range Residual')
        plt.title(f'UKF Residuals Over Time-----{i + 1}')
        plt.ylabel('Range Residual (m)')
        plt.legend()
        plt.grid(True)

        plt.subplot(2, 1, 2)
        plt.plot(time_axis_residuals, residual_data[i, :, 1], linestyle='-', label='Azimuth Residual')
        plt.ylabel('Azimuth Residual (rad)')
        plt.xlabel('Time (s)')
        plt.grid(True)
        plt.legend()
        plt.show()

        plt.figure(figsize=(10, 8))
        time_axis_velocity = np.arange(1, 51) * ST
        plt.subplot(2, 1, 1)
        plt.plot(time_axis_velocity, true_data[i, :, 2], linestyle='-', label='true Velocity')
        plt.plot(time_axis_velocity, tracking_data[i, :, 2], linestyle='--', label='UKF Velocity')
        plt.title('Velocity  & turn_rate Over Time')
        plt.ylabel('Velocity (m/s)')
        plt.xlabel('Time (s)')
        plt.legend()
        plt.grid(True)

        plt.subplot(2, 1, 2)
        plt.plot(time_axis_velocity, true_data[i, :, 4], linestyle='-', label='true turn_rate')
        plt.plot(time_axis_velocity, tracking_data[i, :, 4], linestyle='--', label='UKF turn_rate')
        plt.ylabel('turn_rate (rad/s)')
        plt.xlabel('Time (s)')
        plt.legend()
        plt.grid(True)
        plt.show()


if __name__ == '__main__':
    true_orbit, obs_orbit, tracking_orbit, residuals_history_all = create_ukf_orbit_batch(batch_size=32,
                                                                                          data_len=DATA_LENGTH,
                                                                                          dim_state=DIM_STATE,
                                                                                          dim_obs=DIM_OBS)
    my_plot(true_orbit, tracking_orbit, residuals_history_all)
